import requests
import json
import subprocess
import time
import platform
import psutil
from datetime import datetime

OLLAMA_URL="192.168.1.5"
OLLAMA_MODEL="deekseek-r1:671b"
OLLAMA_KEY="YOUR TOKEN"

# version 0.0.1v
# Warning: test in blank VM only
# 警告：只能在空的虚拟机中测试运行

class OllamaClient:
    def __init__(self, server_ip=OLLAMA_URL, model=OLLAMA_MODEL):
        self.server_url = f"http://{server_ip}:11434/api/generate"
        self.model = model
        self.system_prompt = '''system \"\"\"
        # 指令规则
        如果你认为需要执行指令，请在回复中，单独一行提供，\n执行指令：指令名，例如：\n执行指令：ls, 你认为没有新的执行指令，请你根据你的经验和提交的信息，按照格式给出新的一条指令，用来让系统提供更多的系统信息或者日志，注意不要给出需要交互或者会挂起系统的指令，而是会将信息输出到终端并退出的指令，例如top会进入交互模式会影响通信，而要提供带参数的，例如：top -b -n1，这>样会直接获得信息而不用进入交互模式。
         同时注意，
         (1)不要使用sudo指令，不要使用rm指令，以及任何會修改系統文件的指令
         (2)执行指令要单独一行，以便于客户端识别.
         (3)记住，一定要在回复信息的最后一行附加一条符合执行指令格式的执行指令
        \"\"\"'''
        self.max_retries = 3  # 最大重试次数

    def get_system_info(self):
        """收集完整的系统信息"""
        info = {
            "timestamp": datetime.now().isoformat(),
            "os": f"{platform.system()} {platform.release()}",
            "cpu": {
                "usage": psutil.cpu_percent(),
                "cores": psutil.cpu_count()
            },
            "memory": {
                "total": round(psutil.virtual_memory().total / (1024**3)), 
                "used": psutil.virtual_memory().percent
            },
            "disk": {
                "total": round(psutil.disk_usage('/').total / (1024**3)),
                "used": psutil.disk_usage('/').percent
            },
            "processes": len(psutil.pids())
        }
        return json.dumps(info, indent=2)

    def send_to_ollama(self, prompt):
        """发送请求到Ollama服务"""
        data = {
            "model": self.model,
            "prompt": f"{self.system_prompt}\n{prompt}",
            "stream": False
        }
        
        try:
            response = requests.post(
                self.server_url,
                json=data,
                headers={'Content-Type': 'application/json'},
                timeout=300
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"请求失败: {e}")
            return None

    def extract_command(self, response):
        """严格提取行首的执行指令"""
        if not response or 'response' not in response:
            return None

        # 清理响应文本并逐行检查
        text = response['response'].strip()
        for line in text.split('\n'):
            line = line.strip()
            if line.startswith('执行指令：'):
                return line[len('执行指令：'):].strip()
        return None

    def execute_command(self, command):
        """执行系统命令并返回结果"""
        try:
            result = subprocess.run(
                command,
                shell=True,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            return result.stdout or "命令执行成功（无输出）"
        except subprocess.CalledProcessError as e:
            return f"执行失败: {e.stderr}"

    def interactive_loop(self, interval=10):
        """主交互循环"""
        while True:
            try:
                # 1. 获取系统信息
                sys_info = self.get_system_info()
                print(f"\n[收集系统信息]\n{sys_info}")

                # 2. 初始请求
                current_prompt = f"当前系统状态：\n{sys_info}\n请分析并提供操作指令："
                retry_count = 0

                while retry_count < self.max_retries:
                    # 3. 发送请求并获取响应
                    ai_response = self.send_to_ollama(current_prompt)
                    if not ai_response:
                        print("获取响应失败，等待重试...")
                        time.sleep(interval)
                        continue

                    print(f"\n[AI响应]\n{ai_response['response']}")

                    # 4. 尝试提取指令
                    command = self.extract_command(ai_response)
                    
                    if command:
                        # 5. 执行指令并反馈结果
                        print(f"\n[⚡ 执行指令] {command}")
                        result = self.execute_command(command)
                        print(f"[执行结果]\n{result}")

                        # 将结果作为下一轮输入
                        current_prompt = (
                            f"上次指令执行结果：\n{result}\n"
                            "请分析当前状态并给出下一步指令："
                        )
                        retry_count = 0  # 重置重试计数器
                    else:
                        # 6. 未获取到指令时的追问
                        retry_count += 1
                        current_prompt = (
                            "您未提供有效指令。请必须：\n"
                            "1. 说明需要什么额外信息\n"
                            "2. 在最后一行给出：执行指令：具体命令\n"
                            "当前系统状态：\n" + sys_info
                        )
                        print(f" 未检测到指令，尝试重新获取 ({retry_count}/{self.max_retries})")

                    # 每次交互后等待
                    time.sleep(interval)

                if retry_count >= self.max_retries:
                    print(" 连续多次未获取到有效指令，重启流程...")

            except KeyboardInterrupt:
                print("\n 用户终止操作")
                break
            except Exception as e:
                print(f" 运行时错误: {e}")
                time.sleep(interval * 2)  # 错误时延长等待

if __name__ == "__main__":
    client = OllamaClient()
    print(" Ollama自动化客户端启动 (Ctrl+C退出)")
    client.interactive_loop(interval=15)  # 设置交互间隔为15秒
